import copy
import math
from collections import defaultdict

import numpy as np

from htm.Connection import Connection


class LocationSupeficial2D:

    def __init__(self, params):
        self.cellsPerAxis = params['cellsPerAxis']
        self.scale = params['scale']
        self.rotationMatrix = params['rotationMatrix']
        self.orientation = params['orientation']
        self.cellCoordinateOffsets = params['cellCoordinateOffsets']
        self.anchoringMethod = params['anchoringMethod']
        # Defaults and Calculated Properties
        self.learn = True
        self.cellDimensions = np.asarray([self.cellsPerAxis, self.cellsPerAxis], dtype="int")
        self.moduleMapDimensions = np.asarray([self.scale, self.scale], dtype="float")
        self.phasesPerUnitDistance = 1.0 / self.moduleMapDimensions
        if self.rotationMatrix is None:
            self.rotationMatrix = np.array(
                [[math.cos(self.orientation), -math.sin(self.orientation)],
                 [math.sin(self.orientation), math.cos(self.orientation)]])
            if self.anchoringMethod == "discrete":
                # Need to convert matrix to have integer values
                nonzeros = self.rotationMatrix[np.where(np.abs(self.rotationMatrix) > 0)]
                smallestValue = np.amin(nonzeros)
                self.rotationMatrix /= smallestValue
                self.rotationMatrix = np.ceil(self.rotationMatrix)
        self.connection = Connection(params['connectionParams'])
        # reset every step
        self.activePhases = np.empty((0, 2), dtype="float")
        self.cellsForActivePhases = np.empty(0, dtype="int")
        self.phaseDisplacement = np.empty((0, 2), dtype="float")
        self.activeCells = np.empty(0, dtype="int")
        self.sensoryAssociatedCells = np.empty(0, dtype="int")
        self.activeSegments = np.empty(0, dtype="uint32")

    def reset(self):
        self.activePhases = np.empty((0, 2), dtype="float")
        self.cellsForActivePhases = np.empty(0, dtype="int")
        self.phaseDisplacement = np.empty((0, 2), dtype="float")
        self.activeCells = np.empty(0, dtype="int")
        self.sensoryAssociatedCells = np.empty(0, dtype="int")
        self.activeSegments = np.empty(0, dtype="uint32")

    def getBestSegmentPerCell(self, segmentIds, overlaps):
        maxMatchingSegmentPerCell = defaultdict(int)
        maxPotentialCountOfCell = defaultdict(int)
        for segmentId in segmentIds:
            cell = self.connection.allSegments[segmentId].cell
            count = overlaps[segmentId]
            if count > maxPotentialCountOfCell[segmentId]:
                maxPotentialCountOfCell[cell.cellId] = count
                maxMatchingSegmentPerCell[cell.cellId] = segmentId
        return np.array([maxMatchingSegmentPerCell[key] for key in maxMatchingSegmentPerCell.keys()])

    def sensoryComputeLearningMode(self, anchorInput):
        connectedOverlaps, potentialOverlaps = self.connection.computeActivity(anchorInput)
        activeSegments = self.connection.getActiveSegments(connectedOverlaps)
        matchingSegments = self.connection.getMatchingSegments(potentialOverlaps)
        cellsForActiveSegments = self.connection.mapSegmentsToCells(activeSegments)
        learningActiveSegments = activeSegments[np.in1d(cellsForActiveSegments, self.activeCells)]
        # active cells with no active segments
        remainingCells = np.setdiff1d(self.activeCells, cellsForActiveSegments)
        # active cells with no active segment but with matching segment
        candidateSegmentIds = self.connection.filterSegmentsByCell(matchingSegments, remainingCells)
        cellsForCandidateSegments = self.connection.mapSegmentsToCells(candidateSegmentIds)
        learningMatchingSegments = self.connection.getBestSegmentPerCell(candidateSegmentIds, potentialOverlaps)
        newSegmentCells = np.setdiff1d(remainingCells, cellsForCandidateSegments)
        self.connection.learnOnExisting(learningActiveSegments, anchorInput, anchorInput)
        self.connection.learnOnExisting(learningMatchingSegments, anchorInput, anchorInput)
        self.connection.learnOnNew(newSegmentCells, anchorInput)
        self.activeSegments = activeSegments
        self.sensoryAssociatedCells = self.activeCells

    def sensoryComputeInferenceMode(self, anchorInput):
        if len(anchorInput) == 0:
            return
        connectedOverlaps, _ = self.connection.computeActivity(anchorInput)
        activeSegments = self.connection.getActiveSegments(connectedOverlaps)
        sensorySupportedCells = np.unique(self.connection.mapSegmentsToCells(activeSegments))
        inactivated = np.setdiff1d(self.activeCells, sensorySupportedCells)
        inactivatedIndices = np.array(np.in1d(self.cellsForActivePhases, inactivated)).nonzero()[0]
        if inactivatedIndices.size > 0:
            self.activePhases = np.delete(self.activePhases, inactivatedIndices, axis=0)
        activated = np.setdiff1d(sensorySupportedCells, self.activeCells)
        if "corners" in self.anchoringMethod:
            activatedCoordsBase = np.transpose(np.unravel_index(sensorySupportedCells, self.cellDimensions)).astype('float')
        else:
            activatedCoordsBase = np.transpose(np.unravel_index(activated, self.cellDimensions)).astype('float')
        activatedCoords = np.concatenate(
            [activatedCoordsBase + [iOffset, jOffset] for iOffset in self.cellCoordinateOffsets for jOffset in
             self.cellCoordinateOffsets])
        if "corners" in self.anchoringMethod:
            self.activePhases = activatedCoords / self.cellDimensions
        else:
            if activatedCoords.size > 0:
                self.activePhases = np.append(self.activePhases, activatedCoords / self.cellDimensions, axis=0)
        self.computeActiveCells()
        self.activeSegments = activeSegments
        self.sensoryAssociatedCells = sensorySupportedCells

    def computeActiveCells(self):
        activeCellCoordinates = np.floor(self.activePhases * self.cellDimensions).astype("int")
        self.cellsForActivePhases = (np.ravel_multi_index(activeCellCoordinates.T, self.cellDimensions))
        self.activeCells = np.unique(self.cellsForActivePhases)

    def activateRandomLocation(self):
        self.activePhases = np.array([np.random.random(2)])
        if self.anchoringMethod == "discrete":
            self.activePhases = np.floor(self.activePhases * self.cellDimensions) / self.cellDimensions
        self.computeActiveCells()

    def movementCompute(self, displacement, noiseFactor=0):
        if noiseFactor != 0:
            displacement = copy.deepcopy(displacement)
            xnoise = np.random.normal(0, noiseFactor)
            ynoise = np.random.normal(0, noiseFactor)
            displacement[0] += xnoise
            displacement[1] += ynoise
        phaseDisplacement = (np.matmul(self.rotationMatrix, displacement) * self.phasesPerUnitDistance)
        np.add(self.activePhases, phaseDisplacement, out=self.activePhases)
        np.round(self.activePhases, decimals=9, out=self.activePhases)
        np.mod(self.activePhases, 1.0, out=self.activePhases)
        self.computeActiveCells()
        self.phaseDisplacement = phaseDisplacement

    def sensoryCompute(self, anchorInput, anchorGrowthCandidates):
        if self.learn:
            self.sensoryComputeLearningMode(anchorGrowthCandidates)
        else:
            self.sensoryComputeInferenceMode(anchorInput)

    def getActiveCells(self):
        return self.activeCells

    def getLearnableCells(self):
        return self.activeCells

    def getSensoryAssociatedCells(self):
        return self.sensoryAssociatedCells

    def getNumOfCells(self):
        return np.prod(self.cellDimensions)
